import os
import sys
from CalTorque import CalTorque
from mkplots import MkPlots
from PlotFit import PlotFit_mod
from average_torque_in import AverageIN
from average_in_by_month import MonthlyAverages
from AverageCompareBABreak import CompareBeforeAfterBreak
from AverageCompareSources import CompareSources
import pickle

def SaveAveragePKL(x, y):
    avgdata = {}
    avgdata['x'] = x
    avgdata['y'] = y
    output = open('pickled/pickled_average_in.pkl', 'wb')
    pickle.dump(avgdata, output, -1)
    output.close()

def Plot(dp = 'data/', bins = 400, max_d = 400):
    plot = raw_input('Plot torque vs time and torque vs distance? (y/n)\n>')
    fit = raw_input('Plot Gaussian fit? (y/n)\n>')
    avg = raw_input('Plot average? (y/n) \n>')
    month = raw_input('Plot monthly averages? (y/n) \n>')
    telescope = raw_input('Plot average before and after the telescope break? (y/n)\n>')
    source = raw_input('Plot average for each source type? (y/n)\n>')
    
    data_files = []
    data_files_in = []
    data_files_out = []
    
    ls= os.listdir(dp)
    for i in ls:
        if i.split('.')[-1] not in ['xlsx', 'xls']: continue
        data_files.append(dp+i)
        if 'out' in i:
            data_files_out.append(dp+i)
        elif 'in' in i and 'out' not in i:
            data_files_in.append(dp+i)
    
    if plot == 'y':
        print '\nPlotting torque vs time and torque vs distance'
        MkPlots(data_files)
    else: pass
    
    if fit == 'y':
        print '\nPlotting fits'
        for f in data_files_in:
            print f
            data = CalTorque(f)
            if 'new_parms' not in data.keys(): continue
            PlotFit_mod(data)
    
    if avg == 'y':
        print '\nPlotting average'
        AverageIN(data_files_in, bins, max_d)
    else: pass
    
    if month == 'y':
        print '\nPlotting average by month'
        MonthlyAverages(data_files_in, bins, max_d)
    else: pass
    
    if telescope == 'y':
        print '\nPlotting average before and after telescope break'
        CompareBeforeAfterBreak(data_files_in, bins, max_d)
    else: pass
    
    if source == 'y':
        print '\nPlotting average by source type'
        CompareSources(data_files_in, bins, max_d)
    else: pass

def main(argv = sys.argv):
    Plot()

if __name__ == '__main__':
    main(sys.argv)